feat: store data file size in manifest (#3750)

Part of #3751.

* When we write data files, we store the size in bytes in the manifest.
* We use this size to skip the `HEAD` request needed to find the file
footer byte range.
* When writing a dataset, we can read any missing sizes and fill them
in.

Author: wjones127

java/core/lance-jni/src/file_reader.rs

@@ -92,7 +92,7 @@ fn inner_open<'local>(env: &mut JNIEnv<'local>, file_uri: JString) -> Result<JOb
         let config = SchedulerConfig::max_bandwidth(&obj_store);
         let scan_scheduler = ScanScheduler::new(obj_store, config);
 
-        let file_scheduler = scan_scheduler.open_file(&Path::parse(&path)?).await?;
+        let file_scheduler = scan_scheduler.open_file(&Path::parse(&path)?, None).await?;
         FileReader::try_open(
             file_scheduler,
             None,

java/core/lance-jni/src/fragment.rs

@@ -223,7 +223,7 @@ fn create_fragment<'a>(
 }
 
 const DATA_FILE_CLASS: &str = "com/lancedb/lance/fragment/DataFile";
-const DATA_FILE_CONSTRUCTOR_SIG: &str = "(Ljava/lang/String;[I[III)V";
+const DATA_FILE_CONSTRUCTOR_SIG: &str = "(Ljava/lang/String;[I[IIILjava/lang/Long;)V";
 const DELETE_FILE_CLASS: &str = "com/lancedb/lance/fragment/DeletionFile";
 const DELETE_FILE_CONSTRUCTOR_SIG: &str =
     "(JJLjava/lang/Long;Lcom/lancedb/lance/fragment/DeletionFileType;)V";
@@ -238,6 +238,10 @@ impl IntoJava for &DataFile {
         let path = env.new_string(self.path.clone())?.into();
         let fields = JLance(self.fields.clone()).into_java(env)?;
         let column_indices = JLance(self.column_indices.clone()).into_java(env)?;
+        let file_size_bytes = match self.file_size_bytes {
+            Some(f) => JLance(f as i64).into_java(env)?,
+            None => JObject::null(),
+        };
         Ok(env.new_object(
             DATA_FILE_CLASS,
             DATA_FILE_CONSTRUCTOR_SIG,
@@ -247,6 +251,7 @@ impl IntoJava for &DataFile {
                 JValueGen::Object(&column_indices),
                 JValueGen::Int(self.file_major_version as i32),
                 JValueGen::Int(self.file_minor_version as i32),
+                JValueGen::Object(&file_size_bytes),
             ],
         )?)
     }
@@ -453,12 +458,17 @@ impl FromJObjectWithEnv<DataFile> for JObject<'_> {
         let file_minor_version = env
             .call_method(self, "getFileMinorVersion", "()I", &[])?
             .i()? as u32;
+        let file_size_bytes: Option<i64> = env
+            .call_method(self, "getFileSizeBytes", "()Ljava/lang/Long;", &[])?
+            .l()?
+            .extract_object(env)?;
         Ok(DataFile {
             path,
             fields,
             column_indices,
             file_major_version,
             file_minor_version,
+            file_size_bytes: file_size_bytes.map(|r| r as u64),
         })
     }
 }

java/core/lance-jni/src/traits.rs

@@ -173,6 +173,12 @@ impl IntoJava for JLance<usize> {
     }
 }
 
+impl IntoJava for JLance<i64> {
+    fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
+        Ok(env.new_object("java/lang/Long", "(J)V", &[JValueGen::Long(self.0)])?)
+    }
+}
+
 impl IntoJava for JLance<Option<usize>> {
     fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
         let obj = match self.0 {

java/core/src/main/java/com/lancedb/lance/fragment/DataFile.java

@@ -24,14 +24,21 @@ public class DataFile implements Serializable {
   private final int[] columnIndices;
   private final int fileMajorVersion;
   private final int fileMinorVersion;
+  private final Long fileSizeBytes;
 
   public DataFile(
-      String path, int[] fields, int[] columnIndices, int fileMajorVersion, int fileMinorVersion) {
+      String path,
+      int[] fields,
+      int[] columnIndices,
+      int fileMajorVersion,
+      int fileMinorVersion,
+      Long fileSizeBytes) {
     this.path = path;
     this.fields = fields;
     this.columnIndices = columnIndices;
     this.fileMajorVersion = fileMajorVersion;
     this.fileMinorVersion = fileMinorVersion;
+    this.fileSizeBytes = fileSizeBytes;
   }
 
   public String getPath() {
@@ -54,6 +61,10 @@ public int getFileMinorVersion() {
     return fileMinorVersion;
   }
 
+  public Long getFileSizeBytes() {
+    return fileSizeBytes;
+  }
+
   @Override
   public String toString() {
     return new ToStringBuilder(this)

protos/table.proto

@@ -308,6 +308,13 @@ message DataFile {
   // If both `file_major_version` and `file_minor_version` are set to 0,
   // then this is a version 0.1 or version 0.2 file.
   uint32 file_minor_version = 5;
+
+  // The known size of the file on disk in bytes.
+  //
+  // This is used to quickly find the footer of the file.
+  //
+  // When this is zero, it should be interpreted as "unknown".
+  uint64 file_size_bytes = 6;
 } // DataFile
 
 // Deletion File

python/python/lance/fragment.py

@@ -145,13 +145,16 @@ class DataFile:
         The major version of the data storage format.
     file_minor_version : int
         The minor version of the data storage format.
+    file_size_bytes : Optional[int]
+        The size of the data file in bytes, if available.
     """
 
     _path: str
     fields: List[int]
     column_indices: List[int] = field(default_factory=list)
     file_major_version: int = 0
     file_minor_version: int = 0
+    file_size_bytes: Optional[int] = None
 
     def __init__(
         self,
@@ -160,21 +163,24 @@ def __init__(
         column_indices: List[int] = None,
         file_major_version: int = 0,
         file_minor_version: int = 0,
+        file_size_bytes: Optional[int] = None,
     ):
         # TODO: only we eliminate the path method, we can remove this
         self._path = path
         self.fields = fields
         self.column_indices = column_indices or []
         self.file_major_version = file_major_version
         self.file_minor_version = file_minor_version
+        self.file_size_bytes = file_size_bytes
 
     def __repr__(self):
         # pretend we have a 'path' attribute
         return (
             f"DataFile(path='{self._path}', fields={self.fields}, "
             f"column_indices={self.column_indices}, "
             f"file_major_version={self.file_major_version}, "
-            f"file_minor_version={self.file_minor_version})"
+            f"file_minor_version={self.file_minor_version}, "
+            f"file_size_bytes={self.file_size_bytes})"
         )
 
     @property

python/python/tests/test_fragment.py

@@ -245,7 +245,7 @@ def test_fragment_meta():
     data = {
         "id": 0,
         "files": [
-            {"path": "0.lance", "fields": [0]},
+            {"path": "0.lance", "fields": [0], "file_size_bytes": 100},
             {"path": "1.lance", "fields": [1]},
         ],
         "deletion_file": None,
@@ -261,10 +261,10 @@ def test_fragment_meta():
 
     assert repr(meta) == (
         "FragmentMetadata(id=0, files=[DataFile(path='0.lance', fields=[0], "
-        "column_indices=[], file_major_version=0, file_minor_version=0), "
-        "DataFile(path='1.lance', fields=[1], column_indices=[], "
-        "file_major_version=0, file_minor_version=0)], physical_rows=100, "
-        "deletion_file=None, row_id_meta=None)"
+        "column_indices=[], file_major_version=0, file_minor_version=0, "
+        "file_size_bytes=100), DataFile(path='1.lance', fields=[1], column_indices=[], "
+        "file_major_version=0, file_minor_version=0, file_size_bytes=None)], "
+        "physical_rows=100, deletion_file=None, row_id_meta=None)"
     )
 
 

python/src/file.rs

@@ -398,7 +398,7 @@ impl LanceFileReader {
                 io_buffer_size_bytes: 2 * 1024 * 1024 * 1024,
             },
         );
-        let file = scheduler.open_file(&path).await.infer_error()?;
+        let file = scheduler.open_file(&path, None).await.infer_error()?;
         let inner = FileReader::try_open(
             file,
             None,

python/src/fragment.rs

@@ -658,6 +658,7 @@ impl FromPyObject<'_> for PyLance<DataFile> {
             column_indices: ob.getattr("column_indices")?.extract()?,
             file_major_version: ob.getattr("file_major_version")?.extract()?,
             file_minor_version: ob.getattr("file_minor_version")?.extract()?,
+            file_size_bytes: ob.getattr("file_size_bytes")?.extract()?,
         }))
     }
 }
@@ -679,6 +680,7 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataFile> {
             self.0.column_indices.clone(),
             self.0.file_major_version,
             self.0.file_minor_version,
+            self.0.file_size_bytes,
         ))
     }
 }

rust/lance-file/benches/reader.rs

@@ -66,7 +66,7 @@ fn bench_reader(c: &mut Criterion) {
                         object_store.clone(),
                         SchedulerConfig::default_for_testing(),
                     );
-                    let scheduler = store_scheduler.open_file(file_path).await.unwrap();
+                    let scheduler = store_scheduler.open_file(file_path, None).await.unwrap();
                     let reader = FileReader::try_open(
                         scheduler.clone(),
                         None,
@@ -159,7 +159,7 @@ fn bench_random_access(c: &mut Criterion) {
         let reader = rt.block_on(async move {
             let store_scheduler =
                 ScanScheduler::new(object_store.clone(), SchedulerConfig::default_for_testing());
-            let scheduler = store_scheduler.open_file(file_path).await.unwrap();
+            let scheduler = store_scheduler.open_file(file_path, None).await.unwrap();
             Arc::new(
                 FileReader::try_open(
                     scheduler.clone(),